/*
 * Licensed to the Apache Software Foundation (ASF) under one or more
 * contributor license agreements. See the NOTICE file distributed with
 * this work for additional information regarding copyright ownership.
 * The ASF licenses this file to You under the Apache License, Version 2.0
 * (the "License"); you may not use this file except in compliance with
 * the License. You may obtain a copy of the License at
 *
 *    http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */
package org.apache.kafka.clients.producer.internals;

import java.nio.ByteBuffer;
import java.util.ArrayDeque;
import java.util.Deque;
import java.util.concurrent.TimeUnit;
import java.util.concurrent.locks.Condition;
import java.util.concurrent.locks.ReentrantLock;

import org.apache.kafka.common.MetricName;
import org.apache.kafka.common.errors.TimeoutException;
import org.apache.kafka.common.metrics.Metrics;
import org.apache.kafka.common.metrics.Sensor;
import org.apache.kafka.common.metrics.stats.Meter;
import org.apache.kafka.common.utils.Time;


/**
 * A pool of ByteBuffers kept under a given memory limit. This class is fairly specific to the needs of the producer. In
 * particular it has the following properties:
 * <ol>
 * <li>There is a special "poolable size" and buffers of this size are kept in a free list and recycled
 * <li>It is fair. That is all memory is given to the longest waiting thread until it has sufficient memory. This
 * prevents starvation or deadlock when a thread asks for a large chunk of memory and needs to block until multiple
 * buffers are deallocated.
 * </ol>
 */
public class BufferPool {

    static final String WAIT_TIME_SENSOR_NAME = "bufferpool-wait-time";
    // 内存池可以申请的ByteBuffer字节总大小，由配置文件指定
    private final long totalMemory;
    // free池中ByteBuffer的固定大小；
    private final int poolableSize;
    // 重入锁
    private final ReentrantLock lock;
    /**
     * 空闲的ByteBuffer, 该池子中存储大小等于poolableSize的ByteBuffer；
     * free 对象的使用有点绕，在初始化时，是一个空的Array队列。
     * allocate() 方法是从 free 中取 buffer 或 释放 buffer , deallocate() 是归还 buffer 到 free 中。
     * https://img2020.cnblogs.com/blog/802197/202009/802197-20200907213403925-16683801.png
     *     当 free 空时，从 allocate() 中生产 buffer 对象
     *     deallocate() 方法将 buffer 放到 free 中
     *
     * free 为什么是双向队列
     *     获取 buffer 是从一头取
     *     freeUp() 方法释放 buffer 是从另一头
     * 理论上 allocate() 方法是单线程访问。怕是以防万一吧，一边获取一边释放。
     */
    private final Deque<ByteBuffer> free;
    // 等待分配空间的线程
    private final Deque<Condition> waiters;
    // 内存池中，除了free池和已申请的ByteBuffer，剩余的字节大小。物理上，其处于JVM堆内存中，只是通过nonPooledAvailableMemory标记来约束其可以从堆内存申请的字节大小；
    // ByteBuffer之外的缓冲区，设计为了适应突然的大数据量
    private long nonPooledAvailableMemory;
    private final Metrics metrics;
    private final Time time;
    private final Sensor waitTime;

    // https://img-blog.csdnimg.cn/20201210115209489.png?x-oss-process=image/watermark,type_ZmFuZ3poZW5naGVpdGk,shadow_10,text_aHR0cHM6Ly9ibG9nLmNzZG4ubmV0L2ExMjQwNDY2MTk2,size_16,color_FFFFFF,t_70

    /**
     * 构造方法
     * @param memory The maximum amount of memory that this buffer pool can allocate
     * @param poolableSize The buffer size to cache in the free list rather than deallocating
     * @param metrics instance of Metrics
     * @param time time instance
     * @param metricGrpName logical group name for metrics
     */
    public BufferPool(long memory, int poolableSize, Metrics metrics, Time time, String metricGrpName) {
        this.poolableSize = poolableSize;   // 指定的 poolableSize
        this.lock = new ReentrantLock();    // 初始化 ReentrantLock 锁
        this.free = new ArrayDeque<>();     // 初始化一个 空（empty）的Array队列，存储内存
        this.waiters = new ArrayDeque<>();  // 初始化一个空（empty）的array队列，存储等待线程
        this.totalMemory = memory;  // 总的内存
        this.nonPooledAvailableMemory = memory; // 默认的池外内存，就是总的内存
        // 下面是一些数据统计，不做分析
        this.metrics = metrics;
        this.time = time;
        this.waitTime = this.metrics.sensor(WAIT_TIME_SENSOR_NAME);
        MetricName rateMetricName = metrics.metricName("bufferpool-wait-ratio",  metricGrpName, "The fraction of time an appender waits for space allocation.");
        MetricName totalMetricName = metrics.metricName("bufferpool-wait-time-total", metricGrpName, "The total time an appender waits for space allocation.");
        this.waitTime.add(new Meter(TimeUnit.NANOSECONDS, rateMetricName, totalMetricName));
    }

    /**
     * 分配指定空间的缓存， 如果缓冲区中没有足够的空闲空间，那么会阻塞线程， 直到超时或得到足够空间
     * https://img2020.cnblogs.com/blog/802197/202009/802197-20200907203144906-598326457.png
     *
     * Kafka Producer以ProducerBatch为单位发送数据，而ProducerBatch中的数据以ByteBuffer的形式进行存储。
     *      当发送端数据量极大时，ByteBuffer就会无限制地频繁申请，可能会引发OOM；
     *      另外，发送完数据后，ByteBuffer就会释放，会频繁的引发FullGC，影响Kafka性能。
     *
     * 为此，设计了Kafka Producer端的内存池，它具有以下功能：
     *      限制可以申请的内存总量totalMemory，防止OOM，totalMemory可以通过KafkaProducer配置项指定；
     *      free池持有ProducerRecord的引用，减少FullGC的频率；
     *
     * KafkaProducer发送流程：
     * https://img-blog.csdnimg.cn/20201210090729418.png?x-oss-process=image/watermark,type_ZmFuZ3poZW5naGVpdGk,shadow_10,text_aHR0cHM6Ly9ibG9nLmNzZG4ubmV0L2ExMjQwNDY2MTk2,size_16,color_FFFFFF,t_70
     *
     * @param size The buffer size to allocate in bytes
     * @param maxTimeToBlockMs The maximum time in milliseconds to block for buffer memory to be available
     * @return The buffer
     * @throws InterruptedException If the thread is interrupted while blocked
     * @throws IllegalArgumentException if size is larger than the total memory controlled by the pool (and hence we would block forever)
     */
    public ByteBuffer allocate(int size, long maxTimeToBlockMs) throws InterruptedException {
        // 申请的size大于内存池总大小，抛异常，可以通过KafkaProducer进行重新设置
        if (size > this.totalMemory) {
            throw new IllegalArgumentException("Attempt to allocate " + size + " bytes, but there is a hard limit of " + this.totalMemory + " on memory allocations.");
        }

        ByteBuffer buffer = null;
        // 会有线程争抢，所以需要锁
        this.lock.lock();
        try {
            // 如果有空间大小正合适的空闲buffer， 直接获取并返回
            if (size == poolableSize && !this.free.isEmpty()) {
                return this.free.pollFirst();
            }

            // 判断是否有足够的空闲的内存
            int freeListSize = freeSize() * this.poolableSize;
            if (this.nonPooledAvailableMemory + freeListSize >= size) {
                // 有足够的，未分配的空闲内存
                // 需要整理到一个buffer外空间中，从JVM Heap 中分配内存
                // 如果availableMemory池的容量不够，则释放free池中的ByteBuffer，增加到availableMemory池
                freeUp(size);
                this.nonPooledAvailableMemory -= size;
            } else { // 没有足够空闲的 内存或 buffer
                // 计数器，用于累计已经通过ByteBuffer释放得到的字节大小
                int accumulated = 0;
                // 条件锁, 阻塞自己，等待别的线程释放内存，
                Condition moreMemory = this.lock.newCondition();
                // loop，直到ProducerRecord释放的ByteBuffer足够大
                try {
                    long remainingTimeToBlockNs = TimeUnit.MILLISECONDS.toNanos(maxTimeToBlockMs);
                    // 把自己添加到等待队列中
                    this.waiters.addLast(moreMemory);
                    // 循环 直到有足够空闲，或超时
                    while (accumulated < size) { // 已释放内存 < 要获取的内存 （释放的还不够）
                        // 计时
                        long startWaitNs = time.nanoseconds();
                        long timeNs;
                        boolean waitingTimeElapsed;
                        try {
                            // await超时则退出，返回false；或者被唤醒退出，返回true
                            waitingTimeElapsed = !moreMemory.await(remainingTimeToBlockNs, TimeUnit.NANOSECONDS);
                        } finally {
                            // 还没到最大时长，被唤醒了。更新下已经等待的时长
                            long endWaitNs = time.nanoseconds();
                            timeNs = Math.max(0L, endWaitNs - startWaitNs);
                            this.waitTime.record(timeNs, time.milliseconds());
                        }

                        if (waitingTimeElapsed) {
                            // 等待超时了，不等了。抛出异常，结束
                            throw new TimeoutException("Failed to allocate memory within the configured max blocking time " + maxTimeToBlockMs + " ms.");
                        }

                        remainingTimeToBlockNs -= timeNs;

                        // 是否有释放的刚好足够的空间，否则的话，还得再调整空间
                        // 如果外部释放空间的ByteBuffer大小为poolableSize，则会被放回free池，此时可以从free池获取需要的ByteBuffer
                        if (accumulated == 0 && size == this.poolableSize && !this.free.isEmpty()) {
                            // 有，直接取一个byteBuffer ,返回，结束
                            buffer = this.free.pollFirst();
                            accumulated = size;
                        } else {
                            // 没有足够空闲的，需要调整分配空间 ， 如果分配多了，那么只需要得到 足够size的空间
                            // 例如： 需要 50 ，释放出来了 80 ，那么只取 其中的 50 。
                            // freeUp操作会给availableMemory池释放足够多的字节大小，因为存在两条释放链路，如下:
                            // 当释放的ByteBuffer大小 <= poolableSize：ByteBuffer -> free -> availableMemory
                            // 当释放的ByteBuffer大小 > poolableSize：ByteBuffer -> availableMemory
                            // 详情参考deallocate()方法
                            freeUp(size - accumulated);
                            int got = (int) Math.min(size - accumulated, this.nonPooledAvailableMemory);
                            this.nonPooledAvailableMemory -= got;
                            accumulated += got;
                        }
                    }
                    // Don't reclaim memory on throwable since nothing was thrown
                    accumulated = 0;
                } finally {
                    // 在循环的过程中，有异常了。 那么已经释放出来的空间，再还回去。
                    this.nonPooledAvailableMemory += accumulated;
                    //把自己从等待队列中移除 ， 并结束
                    this.waiters.remove(moreMemory);
                }
            }
        } finally {
            // 后续处理，这里不管分配空间是成功还是失败，都会执行
            try {
                // 三个条件
                // this.nonPooledAvailableMemory == 0 && this.free.isEmpty() : 池外内存为0 ，并且空闲的byteBuffer 没有了。
                // 取反，就是 nonPooledAvailableMemory > 0 || this.free.isNotEmpty() : 池外有内存，或 有空闲的 ByteBuffer
                // !this.waiters.isEmpty() ： 等待队列里有线程正在等待
                if (!(this.nonPooledAvailableMemory == 0 && this.free.isEmpty()) && !this.waiters.isEmpty()) {
                    // 唤醒队列里正在等待的线程
                    this.waiters.peekFirst().signal();
                }
            } finally {
                // 最后的最后，一定得解锁。
                lock.unlock();
            }
        }

        // 到这里，说明空间足够，并且有足够空闲的了。可以执行真正的分配空间了。
        if (buffer == null) {
            // 没有正好的 buffer，从缓冲区外（JVM Heap）中直接分配内存
            // 通过availableMemory池申请ByteBuffer，即在nonPooledAvailableMemory标记足够大的条件下，通过堆内存申请ByteBuffer
            return safeAllocateByteBuffer(size);
        } else {
            // 有正好的 buffer，返回buffer
            return buffer;
        }
    }

    /**
     * Allocate a buffer.  If buffer allocation fails (e.g. because of OOM) then return the size count back to
     * available memory and signal the next waiter if it exists.
     */
    private ByteBuffer safeAllocateByteBuffer(int size) {
        boolean error = true;
        try {
            // 分配空间
            ByteBuffer buffer = allocateByteBuffer(size);
            error = false;
            // 返回buffer
            return buffer;
        } finally {
            if (error) {
                // 分配失败了, 加锁，操作内存pool
                this.lock.lock();
                try {
                    // 归还空间给 池外内存
                    this.nonPooledAvailableMemory += size;
                    if (!this.waiters.isEmpty()) {
                        // 有其他在等待的线程的话，唤醒其他线程
                        this.waiters.peekFirst().signal();
                    }
                } finally {
                    // 加锁不忘解锁
                    this.lock.unlock();
                }
            }
        }
    }

    // 申请内存
    protected ByteBuffer allocateByteBuffer(int size) {
        // 从JVM Heap 中分配空间，并得到持有空间的ByteBuffer对象
        return ByteBuffer.allocate(size);
    }

    /**
     * freeUp操作会给availableMemory池释放足够多的字节大小，因为存在两条释放链路，如下:
     * 	   当释放的ByteBuffer大小 <= poolableSize：ByteBuffer -> free -> availableMemory
     * 	   当释放的ByteBuffer大小 > poolableSize：ByteBuffer -> availableMemory
     * 	   详情参考deallocate()方法
     */
    private void freeUp(int size) {
        while (!this.free.isEmpty() && this.nonPooledAvailableMemory < size) {
            // 循环把 free 里的 byteBuffer 全捞出来，给 nonPooledAvailableMemory
            this.nonPooledAvailableMemory += this.free.pollLast().capacity();
        }
    }

    /**
     * 将缓冲区返回到池中。如果它们是可共享的大小，则将它们添加到空闲列表中，否则只需将内存标记为空闲。
     * 归还 buffer 到 pool 里，即 buffer放回到 free 队列中。 其他的直接标记为 空闲内存就可以了
     * 主要逻辑：
     * 1. 如果 size == poolableSize , 就放到 free 中
     * 2. 如果 size != poolableSize , 归还到 nonPooledAvailableMemory 中. buffer 对象没有引用。等待GC释放
     * 3. 有等待线程的话，唤醒线程
     *
     * @param buffer 要返回的缓冲区
     * @param size 要标记为释放的缓冲区的大小，请注意这可能小于缓冲容量
     *             因为在就地压缩期间缓冲区可能会重新分配自己
     */
    public void deallocate(ByteBuffer buffer, int size) {
        // 照例先加锁
        lock.lock();
        try {
            if (size == this.poolableSize && size == buffer.capacity()) {
                // 如果是完整的buffer，放回到队列里
                buffer.clear();
                this.free.add(buffer);
            } else {
                // 不是完整的buffer，标记为空闲内存就可以了。
                this.nonPooledAvailableMemory += size;
            }
            // 如果有内存的线程，唤醒线程
            Condition moreMem = this.waiters.peekFirst();
            if (moreMem != null) {
                moreMem.signal();
            }
        } finally {
            // 解锁
            lock.unlock();
        }
    }

    public void deallocate(ByteBuffer buffer) {
        deallocate(buffer, buffer.capacity());
    }

    /**
     * the total free memory both unallocated and in the free list
     */
    public long availableMemory() {
        lock.lock();
        try {
            return this.nonPooledAvailableMemory + freeSize() * (long) this.poolableSize;
        } finally {
            lock.unlock();
        }
    }

    // Protected for testing.
    protected int freeSize() {
        return this.free.size();
    }

    /**
     * Get the unallocated memory (not in the free list or in use)
     */
    public long unallocatedMemory() {
        lock.lock();
        try {
            return this.nonPooledAvailableMemory;
        } finally {
            lock.unlock();
        }
    }

    /**
     * 内存中被阻止等待的线程数
     */
    public int queued() {
        lock.lock();
        try {
            return this.waiters.size();
        } finally {
            lock.unlock();
        }
    }

    /**
     * The buffer size that will be retained in the free list after use
     */
    public int poolableSize() {
        return this.poolableSize;
    }

    /**
     * The total memory managed by this pool
     */
    public long totalMemory() {
        return this.totalMemory;
    }

    // package-private method used only for testing
    Deque<Condition> waiters() {
        return this.waiters;
    }
}
